﻿using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ZGSharp.DataSets
{
    public class Mnist
    {
        double[] train_data_double;
        double[] test_data_double;
        int[] train_label_int;
        int[] test_label_int;
        int img_size = 28 * 28;
        public void InitMinst(string filePath)
        {
            byte[] train_data;
            byte[] test_data;
            byte[] train_label;
            byte[] test_label;
            train_data = Decompress(filePath + @"\train-images-idx3-ubyte.gz");
            train_label = Decompress(filePath + @"\train-labels-idx1-ubyte.gz");
            test_data = Decompress(filePath + @"\t10k-images-idx3-ubyte.gz");
            test_label = Decompress(filePath + @"\t10k-labels-idx1-ubyte.gz");
            train_data_double = new double[60000 * img_size];
            train_label_int = new int[60000];
            test_data_double = new double[10000 * img_size];
            test_label_int = new int[10000];
            Parallel.For(0, 60000, i =>
            {
                for (int j = 0; j < img_size; j++)
                {
                    int img_offset = i * img_size + j;
                    train_data_double[img_offset] = train_data[16 + img_offset] / 255.0;
                }
                train_label_int[i] = train_label[8 + i];
            });
            Parallel.For(0, 10000, i =>
            {
                for (int j = 0; j < img_size; j++)
                {
                    int img_offset = i * img_size + j;
                    test_data_double[img_offset] = test_data[16 + img_offset] / 255.0;
                }
                test_label_int[i] = test_label[8 + i];
            });
        }
        public void GetBatchTrainData(int batch_size, double[] data, int[] label)
        {
            Random rdm = new Random();
            Parallel.For(0, batch_size, i => 
            {
                int index = rdm.Next(60000);
                Array.Copy(train_data_double, index * img_size, data, i * img_size, img_size);
                label[i] = train_label_int[index];
            });
        }
        public void GetBatchTestData(int batch_size, double[] data, int[] label)
        {
            Random rdm = new Random();
            Parallel.For(0, batch_size, i =>
            {
                int index = rdm.Next(10000);
                Array.Copy(test_data_double, index * img_size, data, i * img_size, img_size);
                label[i] = test_label_int[index];
            });
        }

        public byte[] Decompress(string fileName)
        {
            FileInfo fileToDecompress = new FileInfo(fileName);
            MemoryStream stream = new MemoryStream();
            using (FileStream originalFileStream = fileToDecompress.OpenRead())
            {

                using (GZipStream decompressionStream = new GZipStream(originalFileStream, CompressionMode.Decompress))
                {
                    decompressionStream.CopyTo(stream);
                }
            }
            return stream.ToArray();
        }
    }
}
